#include <stdio.h>
#include <stdlib.h>
#include <string.h>		
#include <sys/types.h>
#include <arpa/inet.h>		
#include <sys/stat.h>
#include <unistd.h>
#include <netinet/in.h>

#define IPV4_ADDR_LEN		0x0004
#define DNS_REPLY_FLAGS		0x8180
#define DNS_REPLY_REFUSED	0x8183
#define DNS_REPLY_NAME		0xC00C
#define DNS_REPLY_TTL		0x0005
#define DNS_CLASS_IN		0x0001
#define DNS_TYPE_A		    0x0001
#define DNS_TYPE_NS		    0x0002
#define DNS_NUM_ANSWERS		0x0002
#define NAMESERVER_ONE		"ns1"
#define NAMESERVER_TWO		"ns2"
#define WWW			        "www"
#define NS_NAME_ONE		    "\x03ns1\xC0\x0C"
#define NS_NAME_TWO		    "\x03ns2\xC0\x0C"
#define NS_NAME_LEN		    0x0006
#define MAX_DNS_QUESTIONS	1

struct dns_header {
	uint16_t xid;
	uint16_t flags;
	uint16_t num_questions;
	uint16_t num_answers;
	uint16_t num_authority;
	uint16_t num_additional;
};

struct dns_question_section {
	uint16_t type;
	uint16_t class;
};

struct dns_answer_section {
	uint16_t name;
	uint16_t type;
	uint16_t class;
	uint16_t ttl_top;
	uint16_t ttl;
	uint16_t data_len;
};

/* Extract the domain name from the DNS query packet */
char *get_domain_in_question(char *dns_packet, int packet_size)
{
	int dns_header_len = sizeof(struct dns_header);
	if(packet_size > dns_header_len){
        return NULL;
    }
    domain_name_pointer = (dns_packet + dns_header_len);
    
    char *tmp_ptr = NULL;
	char *domain_name = NULL;
	char *domain_name_pointer = NULL;
	int name_part_len = 0;
	int dn_len = 0;
    do {
        /* Get the length of the next part of the domain name */
        name_part_len = (int) domain_name_pointer[0];

        /* If the length is zero or invalid, then stop processing the domain name */
        if((name_part_len <= 0) || (name_part_len > (packet_size - dns_header_len))){
            break;
        }
        domain_name_pointer++;

        /* Reallocate domain_name pointer to name_part_len plus two bytes;
         * one byte for the period, and one more for the trailing NULL byte.
         */
        tmp_ptr = domain_name;
        domain_name = realloc(domain_name,(dn_len + name_part_len + PERIOD_SIZE + 1));
        if(domain_name == NULL){
            if(tmp_ptr) {
                free(tmp_ptr);
            }
            perror("Realloc Failure");
            return NULL;
        }
        memset(domain_name + dn_len, 0, name_part_len + PERIOD_SIZE+1);

        /* Concatenate this part of the domain name, plus the period */
        strncat(domain_name,domain_name_pointer,name_part_len);
        strncat(domain_name,PERIOD,PERIOD_SIZE);

        /* Keep track of how big domain_name is, and point 
         * domain_name_pointer to the next part of the domain name.
         */
        dn_len += name_part_len + PERIOD_SIZE + 1;
        domain_name_pointer += name_part_len;
    } while(name_part_len > 0);

	return domain_name;
}      

int is_dns_request()
{
	int sock = 0, csock = 0, packet_size = 0;

    /* Read in DNS requests */
	struct sockaddr_in clientaddr;
	memset((void *) &clientaddr,0,sizeof(struct sockaddr_in));

    char *dns_packet = NULL;
    if((dns_packet = receive(sock, SOCK_DGRAM, &packet_size, csock, &clientaddr)) == NULL){
        printf("Failed to receive DNS request from client\n");
        return 0;
    }

    /* Process DNS request packets */
    if(packet_size <= (int) (sizeof(struct dns_header) + sizeof(struct dns_question_section))){
        glog("Received invalid DNS packet; packet size too small",LOG_ERROR_TYPE);
        continue;
    }
		
	struct dns_header *header = (struct dns_header *) dns_packet;
    if(ntohs(header->num_questions) != MAX_DNS_QUESTIONS){
        printf("DNS packet contained the wrong number of questions\n");
        return;
    }
		
    /* Extract the domain name in a standard string format */
    char *question_domain = get_domain_in_question(dns_packet,packet_size);

    /* Make sure we got a valid domain query string */
    if(question_domain != NULL && strlen(question_domain) > 0){
        return; 
    }

    /* Check to make sure this is a type A or type NS, class IN DNS query */
    struct dns_question_section *query_info = 
            (struct dns_question_section *) ((dns_packet) + 
            sizeof(struct dns_header) + 
            strlen(question_domain)   + 1);

    if((query_info->class == htons(DNS_CLASS_IN)) && 
      ((query_info->type == htons(DNS_TYPE_A))    || 
       (query_info->type == htons(DNS_TYPE_NS))))
    {
        /* Send DNS reply packet to client */
        if(!send_dns_reply(question_domain, sock, &clientaddr, query_info->type, dns_packet, packet_size)){
            glog("Failed to send DNS response packet",LOG_ERROR_TYPE);
        }
    } 

    if(question_domain) free(question_domain);
	if(dns_packet) free(dns_packet);
	return EXIT_FAILURE;
}     


/* Create DNS reply packet and send it to the client */
int send_dns_reply(char *question_domain, int sock, struct sockaddr_in *clientaddr, 
                   int dns_type, char *request_packet, int request_packet_size)
{
	/* Zero out the answer section structure */
	struct dns_answer_section answer;
	memset(&answer,0,sizeof(struct dns_answer_section));
	int answer_size = sizeof(struct dns_answer_section);

	/* Check to make sure the packet size is of a valid length */
    int hdrlen = sizeof(struct dns_header) + sizeof(struct dns_question_section) + strlen(question_domain);
	if(request_packet_size < hdrlen) { 
        return 0;    
    }
        
    {
    /* Create the DNS answer section */
    answer.name = htons(DNS_REPLY_NAME);
    answer.type = dns_type;
    answer.class = htons(DNS_CLASS_IN);
    answer.ttl = htons(DNS_REPLY_TTL);

	int memcpy_offset = 0;
	int reply_packet_size = 0;
	char *reply_packet = NULL;
	in_addr_t ip_address1 = {0};
	in_addr_t ip_address2 = {0};
    if(dns_type == htons(DNS_TYPE_A)){
        /* Data is an IPv4 address */
        answer.data_len = htons(IPV4_ADDR_LEN);

        /* DNS response packet consists of the original DNS query plus the answer section,
         * plus the answer data (an IPv4 address). We have two IP addresses, so there are
         * two answer sections.
         */
        reply_packet_size = request_packet_size + ((answer_size + IPV4_ADDR_LEN) * DNS_NUM_ANSWERS);
        if((reply_packet = malloc(reply_packet_size)) != NULL) {
            /* Memcpy packet data into the reply packet */
            memcpy(reply_packet,request_packet,request_packet_size);
            memcpy_offset += request_packet_size;
            memcpy(reply_packet+memcpy_offset,(void *) &answer,answer_size);
            memcpy_offset += answer_size;
            memcpy(reply_packet+memcpy_offset,(void *) &ip_address1,IPV4_ADDR_LEN);
            memcpy_offset += IPV4_ADDR_LEN;
            memcpy(reply_packet+memcpy_offset,(void *) &answer,answer_size);
            memcpy_offset += answer_size;
            memcpy(reply_packet+memcpy_offset,(void *) &ip_address2,IPV4_ADDR_LEN);
        } 
        else {
            perror("Malloc Failure");
            return 0;
        }
    } 
    else if(dns_type == htons(DNS_TYPE_NS)){
        answer.data_len = htons(NS_NAME_LEN);

        reply_packet_size = request_packet_size + ((answer_size + NS_NAME_LEN) * DNS_NUM_ANSWERS);
        if((reply_packet = malloc(reply_packet_size)) != NULL){
            /* Memcpy packet data into the reply packet */
            memcpy(reply_packet,request_packet,request_packet_size);
            memcpy_offset += request_packet_size;
            memcpy(reply_packet+memcpy_offset,(void *) &answer,answer_size);
            memcpy_offset += answer_size;
            memcpy(reply_packet+memcpy_offset,NS_NAME_ONE,NS_NAME_LEN);
            memcpy_offset += NS_NAME_LEN;
            memcpy(reply_packet+memcpy_offset,(void *) &answer,answer_size);
            memcpy_offset += answer_size;
            memcpy(reply_packet+memcpy_offset,NS_NAME_TWO,NS_NAME_LEN);
        } 
        else {
            perror("Malloc Failure");
            return 0;
        }
    }

#if 0
    /* Change the number of answers and the flags values of the DNS packet header */
    struct dns_header *header = NULL;
    header = (struct dns_header *) reply_packet;
    header->num_answers = htons(DNS_NUM_ANSWERS);
    header->flags = htons(DNS_REPLY_FLAGS);
#endif        
    
    /* Send reply */
    int bytes_sent = sendto(sock,reply_packet,reply_packet_size,0,(struct sockaddr *) clientaddr, sizeof(struct sockaddr_in));
    if(bytes_sent != reply_packet_size){
        glog("Failed to send response DNS packet",LOG_ERROR_TYPE);
    } 
	return 0;
}       
